import socket
import threading
import sys
import argparse
from datetime import datetime
import ssl
import hashlib
import secrets
import json
from cryptography.fernet import Fernet
from base64 import b64encode, b64decode
import logging

# 设置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class SecureP2PChat:
    def __init__(self, host='0.0.0.0', port=5000):
        # 基本设置
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.host = host
        self.port = port
        self.peer_address = None
        self.username = None
        self.peer_username = None
        self.is_server = False
        
        # 安全相关
        self.max_message_size = 4096  # 最大消息大小
        self.connection_timeout = 30   # 连接超时时间
        self.key = None               # 加密密钥
        self.cipher_suite = None      # 加密套件
        
    def start_server(self):
        """启动监听模式"""
        try:
            self.is_server = True
            self.socket.bind((self.host, self.port))
            self.socket.listen(1)
            logging.info(f"等待连接在 {self.host}:{self.port}...")
            client_socket, address = self.socket.accept()
            self.peer_address = address
            
            # 设置加密
            self.setup_encryption(client_socket)
            
            # 交换用户名
            self.secure_send(client_socket, self.username)
            self.peer_username = self.secure_receive(client_socket)
            
            logging.info(f"{self.peer_username} 已连接！")
            logging.info(f"开始和 {self.peer_username} 聊天吧！(输入quit退出)")
            return client_socket
        except Exception as e:
            logging.error(f"启动失败: {e}")
            sys.exit(1)
            
    def connect_to_peer(self, peer_host, peer_port):
        """连接到对方"""
        try:
            self.is_server = False
            self.socket.connect((peer_host, peer_port))
            self.peer_address = (peer_host, peer_port)
            
            # 设置加密
            self.setup_encryption(self.socket)
            
            # 交换用户名
            self.peer_username = self.secure_receive(self.socket)
            self.secure_send(self.socket, self.username)
            
            logging.info(f"已成功连接到 {self.peer_username}！")
            logging.info(f"开始和 {self.peer_username} 聊天吧！(输入quit退出)")
            return self.socket
        except Exception as e:
            logging.error(f"连接失败: {e}")
            sys.exit(1)
            
    def handle_messages(self, sock, is_server):
        """处理消息收发"""
        # 启动接收消息的线程
        receive_thread = threading.Thread(target=self.receive_messages, args=(sock,))
        receive_thread.daemon = True
        receive_thread.start()
        
        # 发送消息的主循环
        try:
            while True:
                message = input(f"{self.username}> ")
                if message.lower() == 'quit':
                    logging.info("\n正在断开连接...")
                    break
                if message:  # 只发送非空消息
                    formatted_message = f"{self.username}: {message}"
                    self.secure_send(sock, formatted_message)
        except KeyboardInterrupt:
            logging.info("\n正在退出...")
        finally:
            try:
                quit_message = f"\n{self.username} 已退出聊天"
                self.secure_send(sock, quit_message)
            except:
                pass
            sock.close()
            if is_server:
                self.socket.close()
            sys.exit(0)
            
    def receive_messages(self, sock):
        """接收消息的循环"""
        while True:
            try:
                message = self.secure_receive(sock)
                if not message:
                    logging.info(f"\n{self.peer_username} 已断开连接")
                    sock.close()
                    sys.exit(0)
                logging.info(f"\n{message}")
                print(f"{self.username}> ", end='', flush=True)
            except Exception as e:
                logging.error(f"\n连接已断开: {e}")
                sock.close()
                sys.exit(1)

    def generate_session_key(self):
        """生成会话密钥"""
        return Fernet.generate_key()
        
    def setup_encryption(self, sock):
        """建立加密通道"""
        if self.is_server:
            self.key = self.generate_session_key()
            sock.send(self.key)
        else:
            self.key = sock.recv(44)  # Fernet密钥固定长度
        
        self.cipher_suite = Fernet(self.key)
        
    def encrypt_message(self, message):
        """加密消息"""
        return self.cipher_suite.encrypt(message.encode())
        
    def decrypt_message(self, encrypted_message):
        """解密消息"""
        return self.cipher_suite.decrypt(encrypted_message).decode()
        
    def validate_input(self, message):
        """验证输入安全性"""
        if len(message.encode()) > self.max_message_size:
            raise ValueError("消息超过最大长度限制")
        return message.strip()
        
    def secure_send(self, sock, message):
        """安全发送消息"""
        try:
            validated_message = self.validate_input(message)
            encrypted_message = self.encrypt_message(validated_message)
            message_length = len(encrypted_message)
            
            sock.send(message_length.to_bytes(4, 'big'))
            sock.send(encrypted_message)
        except Exception as e:
            logging.error(f"发送失败: {e}")
            raise ConnectionError(f"发送失败: {e}")
            
    def secure_receive(self, sock):
        """安全接收消息"""
        try:
            length_bytes = sock.recv(4)
            if not length_bytes:
                return None
                
            message_length = int.from_bytes(length_bytes, 'big')
            if message_length > self.max_message_size:
                raise ValueError("消息大小超过限制")
                
            encrypted_message = sock.recv(message_length)
            if not encrypted_message:
                return None
                
            return self.decrypt_message(encrypted_message)
        except Exception as e:
            logging.error(f"接收失败: {e}")
            raise ConnectionError(f"接收失败: {e}")

def print_welcome():
    welcome_text = """
    ====================================
    |        Python ChatMini           |
    ====================================
    
    使用说明:
    - 输入'quit'退出程序
    - Ctrl+C也能退出程序  
    - 回车发送消息
    - 空消息不会发送
    
    连接方式:
    - 作为接收方: -p <端口>
    - 作为发起方: IP:端口
    
    v1.0.0
    ====================================
    """
    print(welcome_text)

def main():
    parser = argparse.ArgumentParser(description='点对点聊天程序')
    parser.add_argument('address', nargs='?', help='连接地址 (格式: IP:端口)')
    parser.add_argument('-p', '--port', type=int, default=5000, help='监听端口 (默认: 5000)')
    args = parser.parse_args()
    
    print_welcome()
    chat = SecureP2PChat(port=args.port)
    
    if args.address:
        try:
            host, port = args.address.split(':')
            port = int(port)
            chat.username = input("请输入你的用户名: ")
            logging.info(f"\n正在连接到 {host}:{port}...")
            sock = chat.connect_to_peer(host, port)
            chat.handle_messages(sock, False)
        except ValueError:
            logging.error("地址格式错误。正确格式: IP:端口")
            sys.exit(1)
    else:
        chat.username = input("请输入你的用户名: ")
        sock = chat.start_server()
        chat.handle_messages(sock, True)

if __name__ == "__main__":
    main()